import pandas as pd
import numpy as np
import os


def random_sample_to_excel(input_csv, output_excel, sample_size=100, random_seed=42):
    """
    从CSV文件中随机抽取指定数量的数据，并提取post_id和post_title保存到Excel文件

    参数:
        input_csv: 输入CSV文件路径
        output_excel: 输出Excel文件路径
        sample_size: 需要抽取的样本数量
        random_seed: 随机种子，确保结果可重现
    """
    print(f"从 {input_csv} 中读取数据...")

    try:
        # 读取CSV文件
        df = pd.read_csv(input_csv)
        total_rows = len(df)
        print(f"成功读取数据，共 {total_rows} 行")

        # 确定实际可抽取的样本数
        actual_sample_size = min(sample_size, total_rows)
        print(f"将随机抽取 {actual_sample_size} 行数据")

        # 随机抽样
        np.random.seed(random_seed)
        sampled_df = df.sample(n=actual_sample_size, random_state=random_seed)

        # 只保留post_id和post_title列
        result_df = sampled_df[['post_id', 'post_title']]

        # 保存到Excel
        result_df.to_csv(output_excel, index=False)
        print(f"数据已成功保存到 {output_excel}")

        return result_df

    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        raise


if __name__ == "__main__":
    # 设置文件路径
    input_file = "../data/stock_posts_hk01810_origin_data_cleaned.csv"
    output_file = "../data/random_stock_posts_sample.csv"

    # 确认文件存在
    if not os.path.exists(input_file):
        print(f"错误：找不到输入文件 {input_file}")
        print(f"当前工作目录: {os.getcwd()}")
    else:
        # 抽取1000条数据
        sample_size = 1000
        random_sample_to_excel(input_file, output_file, sample_size)